/*******************************************************************************
 * Copyright 2011 João Antunes
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 ******************************************************************************/
import java.io.*;
import java.util.ArrayList;
import utils.OptionsException;
import utils.Utils;
import aject.attacks.Attack;
import aject.attacks.Test;
import aject.file.FileInterface;
import aject.file.RandomAccessFile;
import aject.protocol.*;

public class Generator {

  public static void printUsage() {
    System.out.println("Usage: generator [tokens payload max_combinations] TPS_FILE OUTPUT_FILE");
    System.out.println();
    System.out.println("Takes some target protocol specification file (TPS_FILE) and generates"
        + "the attacks to be injected by the injector and saves them to "
        + "OUTPUT_FILE. The attacks are simple random generations, or generated by"
        + "combining several elements from the protocol specification and from two files:");
    System.out.println("  Options:");
    System.out.println("    tokens				File with some malicious tokens.");
    System.out.println("    payload				File with some payload tokens.");
    System.out.println("    max_combinations	tokens x tokens x ... x tokens");
    System.out.println();
    System.out.println("Report bugs to <jantunes@di.fc.ul.pt>.");
  }

  public static void main(String[] args) {
    String file_tokens = null;
    String file_payload = null;
    int max_combinations = 0;
    Test test = null;

    String filename_tps = "";
    String filename_output = "";

    // ProtocolSpecification proto_target = new IMAP();
    ProtocolSpecification target_protocol = null;
    try {
      // proto_target = new DebugWords();
      // proto_target = new DebugNumbers();
      // proto_target = new DNS();
      // proto_target = new GeneralProtocolSpecification();
      // proto_target = new XMLProtocolImport("imap.xml",
      // true, // true: simplified
      // new int[] { 0, 1, 4 }, 4);
    } catch (Exception e) {
      Utils.handleException(e, "Error while creating ProtocolSpecification");
    }

    /***************************************************************************
     * check params
     */
    try {

      if (args.length == 2) {
        test = new TestValue();
      } else if (args.length == 5) {
        file_tokens = args[0];
        file_payload = args[1];
        max_combinations = new Integer(args[2]).intValue();
        try {
          test = new TestValue(file_tokens, file_payload, max_combinations);
        } catch (IOException e) {
          throw new OptionsException(OptionsException.Types.MISSING_PARAMETER,
              "Please specify the correct options.");
        }
      } else {
        throw new OptionsException(OptionsException.Types.MISSING_PARAMETER,
            "Please specify the correct options.");
      }

      filename_output = args[args.length - 1];
      filename_tps = args[args.length - 2];

      try {
        File tps_file = new File(filename_tps);
        ObjectInputStream in = new ObjectInputStream(new FileInputStream(tps_file));
        target_protocol = (ProtocolSpecification)in.readObject();
      } catch (Exception e) {
        throw new OptionsException(OptionsException.Types.INCORRECT_PARAMETER,
            "Error while opening TPS_FILE: " + filename_tps);
      }

      try {
        /* copy generated attack to file */
        FileInterface file = new RandomAccessFile(filename_output);
        file.open();
        file.clear();
        System.out
            .println("This process may take some time depending on the type of test and protocol specification.");

        /* header */
        file.writeInt(0);
        file.writeString(target_protocol.getName());
        file.writeString(test.getName());
        file.writeBytes(target_protocol.getPingMessage().toBits().getBytes());

        byte[] attack_message = null; // packet to send
        int total_attacks = 0;

        // For each state in the protocol.
        for (State state : target_protocol.getStates()) {

          // For each packet in that particular state.
          for (MessageSpecification message : state.getMessages()) {
            test.initialize(target_protocol, message);

            // While we can generate new packets to send.
            while ((attack_message = test.generate()) != null) {
              total_attacks++;

              // Print generation completion status.
              if (total_attacks % 50000 == 0)
                System.out.println(" " + total_attacks / 1000 + "k");
              else if (total_attacks % 10000 == 0)
                System.out.print('|');
              else if (total_attacks % 1000 == 0)
                System.out.print('.');

              ArrayList<byte[]> transitions = new ArrayList<byte[]>();

              // Change state of the protocol.
              for (byte[] t : target_protocol.getPath(target_protocol, state))
                transitions.add(t);

              file.writeAttack(new Attack(total_attacks, attack_message, transitions)); // append
                                                                                        // Packet

            } // while
          } // for
        } // for
        System.out.println();
        long pos = file.getPos();
        file.seek(0);
        file.writeInt(total_attacks);
        file.seek(pos);
        file.close();

        System.out
            .println(total_attacks + " generated attacks written to file: " + filename_output);
      } catch (IOException e) {
        throw new OptionsException(OptionsException.Types.INCORRECT_PARAMETER,
            "Error while saving to " + filename_output);
      }

    } catch (OptionsException oe) {
      /* print usage and quit */
      System.out.println(oe.getMessage());
      printUsage();
      System.exit(1);
    }
  }
}
